"""
存储管理器 (Storage Manager)
负责将分类后的文档元素进行结构化存储
支持按类型分离存储和按三级标题分组存储
"""

import json
import os
from typing import List, Dict, Any, Optional
from loguru import logger
from datetime import datetime
import uuid
from pathlib import Path
from .content_classifier import ClassifiedElement


class StorageManager:
    """存储管理器类"""
    
    def __init__(self, base_storage_path: str = "storage"):
        """
        初始化存储管理器
        
        Args:
            base_storage_path: 基础存储路径
        """
        self.base_path = Path(base_storage_path)
        self.session_id = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"
        self.session_path = self.base_path / self.session_id
        
        # 创建存储目录结构
        self._create_storage_structure()
        
    def _create_storage_structure(self):
        """创建存储目录结构"""
        directories = [
            self.session_path,
            self.session_path / "text",
            self.session_path / "text" / "level3_groups",
            self.session_path / "text" / "titles",
            self.session_path / "text" / "narratives",
            self.session_path / "images",
            self.session_path / "images" / "figures",
            self.session_path / "images" / "charts",
            self.session_path / "images" / "diagrams",
            self.session_path / "tables",
            self.session_path / "tables" / "structured",
            self.session_path / "tables" / "data",
            self.session_path / "metadata",
            self.session_path / "raw"
        ]
        
        for directory in directories:
            directory.mkdir(parents=True, exist_ok=True)
            
        logger.info(f"存储目录结构创建完成: {self.session_path}")
    
    def store_classified_elements(self, classified_elements: List[ClassifiedElement], 
                                 level3_groups: Dict[str, List[ClassifiedElement]],
                                 original_file_path: str) -> Dict[str, Any]:
        """
        存储分类后的元素
        
        Args:
            classified_elements: 分类后的元素列表
            level3_groups: 按三级标题分组的元素
            original_file_path: 原始文件路径
            
        Returns:
            存储结果摘要
        """
        logger.info(f"开始存储 {len(classified_elements)} 个分类元素")
        
        storage_result = {
            "session_id": self.session_id,
            "storage_path": str(self.session_path),
            "original_file": original_file_path,
            "timestamp": datetime.now().isoformat(),
            "stored_files": {},
            "statistics": {}
        }
        
        # 1. 按类型分离存储
        type_storage_result = self._store_by_type(classified_elements)
        storage_result["stored_files"].update(type_storage_result)
        
        # 2. 按三级标题分组存储
        level3_storage_result = self._store_level3_groups(level3_groups)
        storage_result["stored_files"].update(level3_storage_result)
        
        # 3. 存储原始数据
        raw_storage_result = self._store_raw_data(classified_elements, original_file_path)
        storage_result["stored_files"].update(raw_storage_result)
        
        # 4. 存储元数据
        metadata_result = self._store_metadata(classified_elements, level3_groups, original_file_path)
        storage_result["stored_files"].update(metadata_result)
        
        # 5. 生成统计信息
        storage_result["statistics"] = self._generate_statistics(classified_elements, level3_groups)
        
        # 6. 保存存储结果摘要
        summary_file = self.session_path / "storage_summary.json"
        with open(summary_file, 'w', encoding='utf-8') as f:
            json.dump(storage_result, f, ensure_ascii=False, indent=2)
        
        logger.info(f"存储完成，会话ID: {self.session_id}")
        return storage_result
    
    def _store_by_type(self, classified_elements: List[ClassifiedElement]) -> Dict[str, str]:
        """
        按类型分离存储元素
        
        Args:
            classified_elements: 分类后的元素列表
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        # 按类型分组
        type_groups = {}
        for element in classified_elements:
            category = element.category
            subcategory = element.subcategory
            
            if category not in type_groups:
                type_groups[category] = {}
            if subcategory not in type_groups[category]:
                type_groups[category][subcategory] = []
            
            type_groups[category][subcategory].append(element)
        
        # 存储文本类型
        if "text" in type_groups:
            text_files = self._store_text_elements(type_groups["text"])
            stored_files.update(text_files)
        
        # 存储图像类型
        if "image" in type_groups:
            image_files = self._store_image_elements(type_groups["image"])
            stored_files.update(image_files)
        
        # 存储表格类型
        if "table" in type_groups:
            table_files = self._store_table_elements(type_groups["table"])
            stored_files.update(table_files)
        
        return stored_files
    
    def _store_text_elements(self, text_groups: Dict[str, List[ClassifiedElement]]) -> Dict[str, str]:
        """
        存储文本元素
        
        Args:
            text_groups: 文本分组
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        for subcategory, elements in text_groups.items():
            # 准备数据
            data = []
            for element in elements:
                data.append({
                    "element_id": element.element_id,
                    "content": element.content,
                    "subcategory": element.subcategory,
                    "confidence": element.confidence,
                    "metadata": element.metadata
                })
            
            # 确定存储路径
            if "title" in subcategory:
                file_path = self.session_path / "text" / "titles" / f"{subcategory}.json"
            else:
                file_path = self.session_path / "text" / "narratives" / f"{subcategory}.json"
            
            # 保存文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            stored_files[f"text_{subcategory}"] = str(file_path)
            logger.debug(f"存储文本类型 {subcategory}: {len(elements)} 个元素 -> {file_path}")
        
        return stored_files
    
    def _store_image_elements(self, image_groups: Dict[str, List[ClassifiedElement]]) -> Dict[str, str]:
        """
        存储图像元素
        
        Args:
            image_groups: 图像分组
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        for subcategory, elements in image_groups.items():
            # 准备数据
            data = []
            for element in elements:
                data.append({
                    "element_id": element.element_id,
                    "content": element.content,  # base64编码的图像数据
                    "subcategory": element.subcategory,
                    "confidence": element.confidence,
                    "metadata": element.metadata
                })
            
            # 确定存储路径
            if subcategory == "chart":
                file_path = self.session_path / "images" / "charts" / f"{subcategory}.json"
            elif subcategory == "diagram":
                file_path = self.session_path / "images" / "diagrams" / f"{subcategory}.json"
            else:
                file_path = self.session_path / "images" / "figures" / f"{subcategory}.json"
            
            # 保存文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            stored_files[f"image_{subcategory}"] = str(file_path)
            logger.debug(f"存储图像类型 {subcategory}: {len(elements)} 个元素 -> {file_path}")
        
        return stored_files
    
    def _store_table_elements(self, table_groups: Dict[str, List[ClassifiedElement]]) -> Dict[str, str]:
        """
        存储表格元素
        
        Args:
            table_groups: 表格分组
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        for subcategory, elements in table_groups.items():
            # 准备数据
            data = []
            for element in elements:
                data.append({
                    "element_id": element.element_id,
                    "content": element.content,
                    "subcategory": element.subcategory,
                    "confidence": element.confidence,
                    "metadata": element.metadata
                })
            
            # 确定存储路径
            if subcategory == "structured":
                file_path = self.session_path / "tables" / "structured" / f"{subcategory}.json"
            else:
                file_path = self.session_path / "tables" / "data" / f"{subcategory}.json"
            
            # 保存文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(data, f, ensure_ascii=False, indent=2)
            
            stored_files[f"table_{subcategory}"] = str(file_path)
            logger.debug(f"存储表格类型 {subcategory}: {len(elements)} 个元素 -> {file_path}")
        
        return stored_files
    
    def _store_level3_groups(self, level3_groups: Dict[str, List[ClassifiedElement]]) -> Dict[str, str]:
        """
        按三级标题分组存储
        
        Args:
            level3_groups: 三级标题分组
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        for level3_number, elements in level3_groups.items():
            # 准备数据
            group_data = {
                "level3_number": level3_number,
                "element_count": len(elements),
                "elements": []
            }
            
            for element in elements:
                group_data["elements"].append({
                    "element_id": element.element_id,
                    "content": element.content,
                    "category": element.category,
                    "subcategory": element.subcategory,
                    "confidence": element.confidence,
                    "metadata": element.metadata
                })
            
            # 生成安全的文件名
            safe_filename = level3_number.replace(".", "_").replace("/", "_")
            file_path = self.session_path / "text" / "level3_groups" / f"group_{safe_filename}.json"
            
            # 保存文件
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(group_data, f, ensure_ascii=False, indent=2)
            
            stored_files[f"level3_group_{level3_number}"] = str(file_path)
            logger.debug(f"存储三级标题组 {level3_number}: {len(elements)} 个元素 -> {file_path}")
        
        return stored_files
    
    def _store_raw_data(self, classified_elements: List[ClassifiedElement], original_file_path: str) -> Dict[str, str]:
        """
        存储原始数据
        
        Args:
            classified_elements: 分类后的元素列表
            original_file_path: 原始文件路径
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        # 存储所有元素的原始数据
        raw_data = {
            "original_file": original_file_path,
            "total_elements": len(classified_elements),
            "elements": []
        }
        
        for element in classified_elements:
            raw_data["elements"].append({
                "element_id": element.element_id,
                "content": element.content,
                "category": element.category,
                "subcategory": element.subcategory,
                "confidence": element.confidence,
                "metadata": element.metadata
            })
        
        file_path = self.session_path / "raw" / "all_elements.json"
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(raw_data, f, ensure_ascii=False, indent=2)
        
        stored_files["raw_data"] = str(file_path)
        logger.debug(f"存储原始数据: {len(classified_elements)} 个元素 -> {file_path}")
        
        return stored_files
    
    def _store_metadata(self, classified_elements: List[ClassifiedElement], 
                       level3_groups: Dict[str, List[ClassifiedElement]], 
                       original_file_path: str) -> Dict[str, str]:
        """
        存储元数据
        
        Args:
            classified_elements: 分类后的元素列表
            level3_groups: 三级标题分组
            original_file_path: 原始文件路径
            
        Returns:
            存储文件路径字典
        """
        stored_files = {}
        
        # 生成元数据
        metadata = {
            "session_info": {
                "session_id": self.session_id,
                "timestamp": datetime.now().isoformat(),
                "original_file": original_file_path
            },
            "processing_info": {
                "total_elements": len(classified_elements),
                "level3_groups_count": len(level3_groups),
                "processing_stage": "stage1_deconstruction"
            },
            "element_distribution": {},
            "level3_distribution": {}
        }
        
        # 统计元素分布
        for element in classified_elements:
            category_key = f"{element.category}.{element.subcategory}"
            if category_key not in metadata["element_distribution"]:
                metadata["element_distribution"][category_key] = 0
            metadata["element_distribution"][category_key] += 1
        
        # 统计三级标题分布
        for level3_number, elements in level3_groups.items():
            metadata["level3_distribution"][level3_number] = len(elements)
        
        file_path = self.session_path / "metadata" / "session_metadata.json"
        with open(file_path, 'w', encoding='utf-8') as f:
            json.dump(metadata, f, ensure_ascii=False, indent=2)
        
        stored_files["metadata"] = str(file_path)
        logger.debug(f"存储元数据 -> {file_path}")
        
        return stored_files
    
    def _generate_statistics(self, classified_elements: List[ClassifiedElement], 
                           level3_groups: Dict[str, List[ClassifiedElement]]) -> Dict[str, Any]:
        """
        生成统计信息
        
        Args:
            classified_elements: 分类后的元素列表
            level3_groups: 三级标题分组
            
        Returns:
            统计信息字典
        """
        statistics = {
            "total_elements": len(classified_elements),
            "categories": {},
            "subcategories": {},
            "level3_groups": {
                "count": len(level3_groups),
                "groups": {}
            },
            "confidence_stats": {
                "average": 0.0,
                "high_confidence": 0,  # >= 0.8
                "medium_confidence": 0,  # 0.6-0.8
                "low_confidence": 0  # < 0.6
            }
        }
        
        # 统计类别和子类别
        confidence_sum = 0.0
        for element in classified_elements:
            # 类别统计
            if element.category not in statistics["categories"]:
                statistics["categories"][element.category] = 0
            statistics["categories"][element.category] += 1
            
            # 子类别统计
            subcategory_key = f"{element.category}.{element.subcategory}"
            if subcategory_key not in statistics["subcategories"]:
                statistics["subcategories"][subcategory_key] = 0
            statistics["subcategories"][subcategory_key] += 1
            
            # 置信度统计
            confidence_sum += element.confidence
            if element.confidence >= 0.8:
                statistics["confidence_stats"]["high_confidence"] += 1
            elif element.confidence >= 0.6:
                statistics["confidence_stats"]["medium_confidence"] += 1
            else:
                statistics["confidence_stats"]["low_confidence"] += 1
        
        # 计算平均置信度
        if len(classified_elements) > 0:
            statistics["confidence_stats"]["average"] = confidence_sum / len(classified_elements)
        
        # 统计三级标题组
        for level3_number, elements in level3_groups.items():
            statistics["level3_groups"]["groups"][level3_number] = {
                "element_count": len(elements),
                "categories": {}
            }
            
            # 统计每个组内的类别分布
            for element in elements:
                category = element.category
                if category not in statistics["level3_groups"]["groups"][level3_number]["categories"]:
                    statistics["level3_groups"]["groups"][level3_number]["categories"][category] = 0
                statistics["level3_groups"]["groups"][level3_number]["categories"][category] += 1
        
        return statistics
    
    def get_session_path(self) -> str:
        """获取当前会话的存储路径"""
        return str(self.session_path)
    
    def get_session_id(self) -> str:
        """获取当前会话ID"""
        return self.session_id